#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
增强版AI标题生成器
支持分批处理（每批最多30个产品）和中文标题生成
"""

import os
import json
import base64
import re
import time
import pandas as pd
from volcenginesdkarkruntime import Ark


def build_product_key(merchant_info, source_info):
    """根据商家与货源信息生成唯一的目录键"""
    merchant = str(merchant_info or "").strip()
    source = str(source_info or "").strip()
    if not merchant and not source:
        return None
    return f"{merchant}-{source}"


def format_row_number(idx):
    """格式化行号显示"""
    return idx + 1 if isinstance(idx, int) else idx

def load_api_config():
    """加载API配置"""
    try:
        with open('api_key.json', 'r', encoding='utf-8') as f:
            config = json.load(f)
        return config
    except Exception as e:
        print(f"❌ 加载API配置失败: {e}")
        return None

def encode_image_to_base64(image_path):
    """将图片编码为base64"""
    try:
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
    except Exception as e:
        print(f"❌ 图片编码失败 {image_path}: {e}")
        return None

def read_excel_data():
    """读取Excel数据"""
    try:
        excel_file = "production.xlsx"
        if not os.path.exists(excel_file):
            print(f"❌ Excel文件不存在: {excel_file}")
            return None
        
        df = pd.read_excel(excel_file)
        print(f"📖 成功读取Excel文件: {excel_file}")
        print(f"📋 Excel文件包含 {len(df)} 行数据")
        
        return df
    except Exception as e:
        print(f"❌ 读取Excel文件失败: {e}")
        return None

def get_product_image_paths(df):
    """根据Excel数据获取产品图片路径及对应唯一键"""
    image_infos = []
    downloads_dir = "downloads"

    if not os.path.exists(downloads_dir):
        print(f"❌ downloads目录不存在: {downloads_dir}")
        return []

    print(f"🔍 开始查找产品图片...")

    for idx, row in df.iterrows():
        merchant_info = row.get('商家信息', '')
        source_info = row.get('货源信息', '')
        product_key = build_product_key(merchant_info, source_info)

        info = {
            "index": idx,
            "key": product_key,
            "merchant": str(merchant_info or "").strip(),
            "source": str(source_info or "").strip(),
            "path": None,
        }

        display_idx = format_row_number(idx)

        if not product_key:
            print(f"⚠️ 第{display_idx}行缺少商家信息或货源信息，跳过图片匹配")
            image_infos.append(info)
            continue

        product_dir = os.path.join(downloads_dir, product_key)
        if not os.path.exists(product_dir):
            print(f"⚠️ 第{display_idx}行对应的图片目录不存在: {product_dir}")
            image_infos.append(info)
            continue

        found_image = None
        for root, _, files in os.walk(product_dir):
            for file in files:
                if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                    found_image = os.path.join(root, file)
                    break
            if found_image:
                break

        if found_image:
            info["path"] = found_image
            print(f"✅ 第{display_idx}行找到图片: {os.path.basename(found_image)}")
        else:
            print(f"⚠️ 第{display_idx}行未找到图片文件")

        image_infos.append(info)

    valid_images = [info for info in image_infos if info["path"]]
    print(f"📊 总共找到 {len(valid_images)}/{len(image_infos)} 张有效图片")

    return image_infos

def generate_titles_with_ai_batch(image_infos, api_config, batch_size=30):
    """使用AI分批生成商品标题，返回索引与唯一键的映射"""
    try:
        print("🤖 正在使用AI分批生成商品标题...")

        valid_image_data = [info for info in image_infos if info.get("path")]

        if not valid_image_data:
            print("❌ 没有可用的图片进行AI分析")
            return {"by_index": {}, "by_key": {}}

        total_images = len(valid_image_data)
        print(f"📸 总共需要处理 {total_images} 张图片")
        print(f"📦 分批处理，每批最多 {batch_size} 张图片")

        # 初始化客户端
        client = Ark(
            base_url=api_config['base_url'],
            api_key=api_config['api_key']
        )

        titles_by_index = {}
        titles_by_key = {}

        # 分批处理
        for batch_start in range(0, total_images, batch_size):
            batch_end = min(batch_start + batch_size, total_images)
            batch_data = valid_image_data[batch_start:batch_end]
            batch_num = (batch_start // batch_size) + 1
            total_batches = (total_images + batch_size - 1) // batch_size

            print(f"\n📦 处理第 {batch_num}/{total_batches} 批 ({len(batch_data)} 张图片)")

            # 准备当前批次的图片内容
            image_contents = []
            batch_entries = []

            for entry in batch_data:
                base64_image = encode_image_to_base64(entry["path"])
                if base64_image:
                    image_contents.append({
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{base64_image}"
                        }
                    })
                    batch_entries.append(entry)
                    print(f"   ✅ 已编码图片: {os.path.basename(entry['path'])}")
                else:
                    print(f"   ❌ 图片编码失败: {os.path.basename(entry['path'])}")

            if not image_contents:
                print(f"   ⚠️ 第 {batch_num} 批没有成功编码的图片，跳过")
                continue

            # 构建消息（使用新的提示词）
            messages = [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": (
                                "我是电商从业者，请为我发送的以下男士休闲运动鞋图片按顺序逐一编写商品标题。"
                                "要求：每张图片生成一个对应标题，严格按照给出图片的顺序依次输出。"
                                "输出格式为\"序号. 标题\"，例如\"1. 男士休闲运动鞋...\"。"
                                "如果某张图片无法生成合适标题，请输出\"序号.\"后直接留空。"
                                "标题需为中文，不含品牌名称，不能提及图片中的水印内容，不能包含'真皮'、'皮'、'麂皮'等皮类词语。"
                            )
                        }
                    ] + image_contents
                }
            ]

            # 调用API
            print(f"   🔄 正在调用AI API (第 {batch_num} 批)...")
            try:
                response = client.chat.completions.create(
                    model=api_config['model'],
                    messages=messages,
                    max_tokens=3000,  # 增加token数量以支持更多标题
                    temperature=0.7,
                    extra_headers={'x-is-encrypted': 'true'}
                )
                
                # 解析响应
                if response.choices and len(response.choices) > 0:
                    ai_response = response.choices[0].message.content
                    print(f"   ✅ AI响应成功 (第 {batch_num} 批)")
                    print(f"   📝 AI原始响应:")
                    print(f"   {ai_response}")

                    # 解析当前批次的标题
                    batch_titles = parse_ai_titles_chinese(ai_response, len(image_contents))

                    for entry, title in zip(batch_entries, batch_titles):
                        original_index = entry["index"]
                        product_key = entry.get("key")
                        display_idx = format_row_number(original_index)
                        titles_by_index[original_index] = title
                        if product_key:
                            titles_by_key[product_key] = title
                        if title:
                            print(f"   ✅ 标题 匹配行{display_idx}({product_key or '索引匹配'}): {title}")
                        else:
                            print(f"   ⚠️ 标题 匹配行{display_idx}({product_key or '索引匹配'}): 空标题")

                else:
                    print(f"   ❌ 第 {batch_num} 批AI未返回有效响应")

            except Exception as e:
                print(f"   ❌ 第 {batch_num} 批API调用失败: {e}")
                continue

            # 批次间等待，避免API限制
            if batch_num < total_batches:
                print(f"   ⏳ 等待2秒后处理下一批...")
                time.sleep(2)

        print(f"\n📊 分批处理完成，共生成 {len(titles_by_index)} 个标题")
        return {"by_index": titles_by_index, "by_key": titles_by_key}

    except Exception as e:
        print(f"❌ AI分批生成标题失败: {e}")
        return {"by_index": {}, "by_key": {}}

def parse_ai_titles_chinese(ai_response, expected_count):
    """解析AI响应中的中文标题"""
    titles = []
    lines = ai_response.strip().split('\n')
    
    print(f"\n🔍 解析AI响应中的中文标题...")
    
    # 中文禁用词列表
    banned_keywords = ['真皮', '皮革', '麂皮', '牛皮', '羊皮', '皮质']
    # 中文必需关键词（至少包含一个）
    required_keywords = ['男', '鞋', '运动鞋', '休闲鞋', '板鞋', '跑鞋', '篮球鞋', '训练鞋']

    def normalize_title_text(raw_text: str, idx: int) -> str:
        text = raw_text.strip()
        text = re.sub(r'^[-*•]+\s*', '', text)
        text = text.strip()
        if not text:
            print(f"   ⚠️ 标题 {idx}: 内容为空，保持空标题")
            return ""

        # 检查是否包含中文字符
        has_chinese = bool(re.search(r'[\u4e00-\u9fff]', text))
        contains_required = any(keyword in text for keyword in required_keywords)
        contains_banned = any(word in text for word in banned_keywords)
        is_long_enough = len(text) >= 5

        if not has_chinese:
            print(f"   ⚠️ 标题 {idx}: 不包含中文字符，置为空: {text}")
            return ""

        reasons = []
        if not contains_required:
            reasons.append("缺少鞋类关键词")
        if contains_banned:
            reasons.append("包含禁用词")
        if not is_long_enough:
            reasons.append("长度不足")

        if reasons:
            print(f"   ⚠️ 标题 {idx}: {'、'.join(reasons)}，置为空: {text}")
            return ""

        print(f"   ✅ 提取中文标题 {idx}: {text}")
        return text

    numbered_entries = {}
    plain_entries = []

    for raw_line in lines:
        stripped = raw_line.strip()
        if not stripped:
            continue

        match = re.match(r'^(\d+)[\\.、·-]?\s*(.*)$', stripped)
        if match:
            idx = int(match.group(1))
            content = match.group(2)
            numbered_entries[idx] = content
        else:
            plain_entries.append(stripped)

    print(f"\n🔍 开始按序号解析标题（期望 {expected_count} 个）...")

    for idx in range(1, expected_count + 1):
        if idx in numbered_entries:
            raw_title = numbered_entries[idx]
        elif plain_entries:
            raw_title = plain_entries.pop(0)
            print(f"   ℹ️ 标题 {idx}: AI 未提供序号，按顺序匹配 -> {raw_title}")
        else:
            raw_title = ""
        titles.append(normalize_title_text(raw_title, idx))

    valid_count = sum(1 for t in titles if t)
    print(f"\n📊 成功解析 {valid_count}/{expected_count} 个有效中文标题")
    
    return titles

def update_excel_with_titles_batch(df, image_infos, titles_mapping):
    """使用分批生成的标题更新Excel（优先按目录键匹配，其次按索引回退）"""
    try:
        print(f"\n📝 开始更新Excel中的产品标题...")

        # 确定标题列名
        title_column = None
        for candidate in ['产品标题', '产品标题:']:
            if candidate in df.columns:
                title_column = candidate
                break

        if title_column is None:
            title_column = '产品标题'
            df[title_column] = None
            print(f"⚠️ 未找到现有标题列，已创建新列: {title_column}")
        else:
            print(f"✅ 使用标题列: {title_column}")

        # 确保列类型为对象，避免写入字符串时的兼容性警告
        df[title_column] = df[title_column].astype("object")

        titles_by_index = titles_mapping.get('by_index', {})
        titles_by_key = titles_mapping.get('by_key', {})

        updated_count = 0

        for info in image_infos:
            idx = info['index']
            product_key = info.get('key')
            image_path = info.get('path')
            display_idx = format_row_number(idx)

            if image_path is None:
                print(f"   ⚠️ 第{display_idx}行: 无图片或目录缺失，保持原标题")
                continue

            if idx not in df.index:
                print(f"   ⚠️ 第{display_idx}行: 索引超出当前DataFrame范围")
                continue

            new_title = None
            match_strategy = None

            if product_key and product_key in titles_by_key:
                new_title = titles_by_key[product_key]
                match_strategy = "目录键匹配"
            elif idx in titles_by_index:
                new_title = titles_by_index[idx]
                match_strategy = "索引回退"
            else:
                print(f"   ⚠️ 第{display_idx}行: 未找到匹配的AI标题（key={product_key}）")
                continue

            old_title = df.at[idx, title_column] if title_column in df.columns else ''
            df.at[idx, title_column] = new_title
            print(f"   ✅ 第{display_idx}行({match_strategy}): {old_title} → {new_title}")

            updated_count += 1

        # 保存Excel文件
        excel_file = "production.xlsx"
        df.to_excel(excel_file, index=False)
        print(f"\n✅ Excel文件已更新: {excel_file}")
        print(f"📊 共更新了 {updated_count} 个产品标题")
        
        return True
        
    except Exception as e:
        print(f"❌ 更新Excel失败: {e}")
        return False

def main():
    """主函数"""
    print("🚀 增强版AI标题生成器启动")
    print("="*50)
    
    # 1. 加载API配置
    print("1️⃣ 加载API配置...")
    api_config = load_api_config()
    if not api_config:
        print("❌ API配置加载失败，程序退出")
        return
    
    print(f"✅ API配置加载成功")
    print(f"   模型: {api_config['model']}")
    print(f"   API地址: {api_config['base_url']}")
    
    # 2. 读取Excel数据
    print("\n2️⃣ 读取Excel数据...")
    df = read_excel_data()
    if df is None:
        print("❌ Excel数据读取失败，程序退出")
        return
    
    # 3. 获取产品图片路径
    print("\n3️⃣ 获取产品图片路径...")
    image_infos = get_product_image_paths(df)
    if not any(info.get('path') for info in image_infos):
        print("❌ 没有找到任何产品图片，程序退出")
        return
    
    # 4. 使用AI分批生成标题
    print("\n4️⃣ 使用AI分批生成中文标题...")
    titles_dict = generate_titles_with_ai_batch(image_infos, api_config, batch_size=30)
    if not titles_dict.get('by_index'):
        print("❌ AI标题生成失败，程序退出")
        return
    
    # 5. 更新Excel
    print("\n5️⃣ 更新Excel...")
    success = update_excel_with_titles_batch(df, image_infos, titles_dict)
    
    if success:
        print("\n🎉 增强版AI标题生成完成！")
        print(f"📊 处理统计:")
        print(f"   - 总产品数: {len(df)}")
        print(f"   - 有图片产品数: {len([info for info in image_infos if info.get('path')])}")
        print(f"   - 成功生成标题数: {len(titles_dict.get('by_index', {}))}")
    else:
        print("\n❌ 程序执行失败")

if __name__ == "__main__":
    main()
